import re
import os
import math
import queue
import threading
from typing import Iterable

from .config import get_option
from .lang import convert_japanese_romaji, fix_chinese_english_spacing, detect_language

_g_piper_engines = {}
_g_piper_engines_lock = threading.Lock()

def split_at_puncation_marks(text: str) -> list[str]:
    # 定义全角标点和对应的半角标点
    punctuations = '，、。！？；：“”‘’（）【】{}《》～,,!?;:""\'\'()[]{}<>~'
    lines = re.split('[\\n' + re.escape(punctuations) + ']|\\.\\s', text)
    # lines = re.split(r'[\n！？。]|[!\?\.]\s', text)
    return [line for line in (line.strip() for line in lines) if line]

def convert_text_readable(text: str) -> list[str]:
    md_link_pattern = r'\[([^\]]+)\]\([^\)]+\)'
    url_pattern = r'https?://([a-zA-Z0-9\-_\.]+)(:[0-9]+)?(/[a-zA-Z0-9\-_\./%]*)?'
    code_block_pattern = r'```(\w*)\n(\n|.)*?\n```'
    text = re.sub(url_pattern, r'\1', text)
    text = re.sub(md_link_pattern, r'\1', text)
    text = re.sub(code_block_pattern, r'[\1 code block]', text)
    text = text.replace('**', '')
    text = text.replace('tten', 'ttin')  # bug of piper?
    text = convert_japanese_romaji(text, mode='hanzi')
    text = fix_chinese_english_spacing(text)
    return split_at_puncation_marks(text)

download_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'models')

def get_lang_engine(lang: str, quality: str = 'low'):
    try:
        import piper
    except ImportError:
        return None
    if lang not in _g_piper_engines:
        with _g_piper_engines_lock:
            if lang not in _g_piper_engines:
                path_lut = {
                    'medium': {
                        'zh': os.path.join(download_path, 'zh_CN-huayan-medium.onnx'),
                        'en': os.path.join(download_path, 'en_US-hfc_female-medium.onnx'),
                    },
                    'low': {
                        'zh': os.path.join(download_path, 'zh_CN-huayan-x_low.onnx'),
                        'en': os.path.join(download_path, 'en_US-danny-low.onnx'),
                    },
                }
                if quality not in path_lut or lang not in path_lut[quality]:
                    engine = None
                else:
                    path = path_lut[quality][lang]
                    cfg_path = path + '.json'
                    if not os.path.isfile(path) or not os.path.isfile(cfg_path):
                        engine = None
                    else:
                        engine = piper.PiperVoice.load(path, use_cuda=False)
                _g_piper_engines[lang] = engine
    return _g_piper_engines[lang]

def speech_generator(lines: Iterable[str], speed: float, silence: float, rate: int, quality: str) -> Iterable[bytes]:
    lang = 'zh'
    for line in lines:
        line = line.strip()
        if not line:
            continue
        lang = detect_language(line, lang)
        engine = get_lang_engine(lang, quality)
        if engine:
            stream = engine.synthesize_stream_raw(line, length_scale=1 / speed, sentence_silence=silence)
            if rate == engine.config.sample_rate:
                yield from stream
            else:
                import audioop
                state = None
                for chunk in stream:
                    chunk, state = audioop.ratecv(chunk, 2, 1, engine.config.sample_rate, rate, state)
                    yield chunk

def adjust_volume(gen: Iterable[bytes], volume: float) -> Iterable[bytes]:
    import audioop
    for chunk in gen:
        chunk = audioop.mul(chunk, 2, volume)
        yield chunk

def adjust_pitch(gen: Iterable[bytes], pitch: float, rate: int) -> Iterable[bytes]:
    import audioop
    state = None
    pitch = math.exp2(pitch / 5)
    for chunk in gen:
        chunk, state = audioop.ratecv(chunk, 2, 1, rate, round(rate / pitch), state)
        yield chunk

def text_to_speech(text: str | Iterable[str], speed: float = 1.0, volume: float = 1.0, pitch: float = 0.0, silence: float = 0.0, rate: int = 0, quality: str = 'medium') -> tuple[Iterable[bytes], int]:
    if isinstance(text, str):
        lines = convert_text_readable(text)
    else:
        lines = text
    volume = min(1, max(0, volume))
    pitch = min(3, max(-3, pitch))
    speed /= math.exp2(pitch / 5)
    speed = min(4, max(0.25, speed))
    silence = min(2, max(0, silence))
    if rate == 0:
        if quality == 'low':
            rate = 16000
        else:
            rate = 22050
    if quality == 'high':
        quality = 'medium'  # high quality not available yet
    rate = min(22050, max(8000, rate))
    gen = speech_generator(lines, speed, silence, rate, quality)
    if volume != 1:
        gen = adjust_volume(gen, volume)
    if pitch != 0:
        gen = adjust_pitch(gen, pitch, rate)
    return gen, rate

def ffplay_speech(audio: bytes, rate: int):
    import tempfile
    import subprocess
    import wave
    with tempfile.NamedTemporaryFile(suffix='.wav') as file:
        with wave.open(file, 'wb') as wavfile:
            wavfile.setnchannels(1)
            wavfile.setsampwidth(2)
            wavfile.setframerate(rate)
            wavfile.writeframes(audio)
        file.flush()
        subprocess.check_call(['ffplay', '-loglevel', 'error', '-nodisp', '-autoexit', file.name], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

_g_queue_text = queue.Queue()
_g_queue_audio = queue.Queue()

def synth_text_thread():
    progress = 0
    while True:
        text, finish, options = _g_queue_text.get()
        lines = convert_text_readable(text)
        if len(lines) > progress + 1:
            if progress == 0:
                    clear_pending_audio()
            text = lines[progress:len(lines) - 1]
            stream, rate = text_to_speech(text, **options)
            audio = b''.join(stream)
            progress = len(lines) - 1
            _g_queue_audio.put((audio, rate))
        if finish:
            text = lines[progress:]
            stream, rate = text_to_speech(text, **options)
            audio = b''.join(stream)
            # （✿◡‿◡）
            # （´｡• ᵕ •｡`) ♡
            # （づ￣ 3￣)づ 💕
            progress = 0
            _g_queue_audio.put((audio, rate))

def play_text_thread():
    import pyaudio
    pa = pyaudio.PyAudio()
    stream = None
    while True:
        audio, rate = _g_queue_audio.get()
        if not stream:
            stream = pa.open(format=pyaudio.paInt16, channels=1, rate=rate, output=True)
            assert stream
        stream.write(audio)
        stream.write(b'\0' * round(rate * 0.1))

_g_play_text_thread = None
_g_synth_text_thread = None

def clear_pending_audio():
    if not _g_queue_audio.empty():
        while not _g_queue_audio.empty():
            try:
                _g_queue_audio.get_nowait()
            except queue.Empty:
                break

def play_audio_text(text: str, finish: bool):
    options = get_option('tts_options', {})
    if options.get('volume', 1) <= 0:
        return
    global _g_play_text_thread
    global _g_synth_text_thread
    if not _g_play_text_thread or not _g_play_text_thread.is_alive():
        _g_play_text_thread = threading.Thread(target=play_text_thread, daemon=True)
        _g_play_text_thread.start()
    if not _g_synth_text_thread or not _g_synth_text_thread.is_alive():
        _g_synth_text_thread = threading.Thread(target=synth_text_thread, daemon=True)
        _g_synth_text_thread.start()
    _g_queue_text.put((text, finish, options))

if __name__ == '__main__':
    import time
    t0 = time.time()
    play_audio_text('I have written a Python implementation of Radix Sort', False)
    play_audio_text('I have written a Python implementation of Radix Sort for you. The code snippet is displayed', False)
    play_audio_text('I have written a Python implementation of Radix Sort for you. The code snippet is displayed in a separate window.', False)
    play_audio_text('I have written a Python implementation of Radix Sort for you. The code snippet is displayed in a separate window. If you need any further assistance or', False)
    play_audio_text('I have written a Python implementation of Radix Sort for you. The code snippet is displayed in a separate window. If you need any further assistance or explanation, feel free to ask!', True)
    time.sleep(9)
